#pragma comment(linker,"/nodefaultlib:libc")
#pragma comment(linker,"/nodefaultlib:msvcrt")
#pragma comment(linker,"/entry:DllEntry")
#pragma comment(linker,"/fixed")
#pragma comment(linker,"/align:4096")
#pragma comment(linker,"/base:0x70010000")

#define WIN32_LEAN_AND_MEAN		// Exclude rarely-used stuff from Windows headers
#define _WIN32_WINNT 0x0501
#include <windows.h>
#include <WINSOCK2.H>
#include <WinDNS.h>
#include <stdio.h>
#include <malloc.h>
#pragma comment(lib,"ntdll.lib")
#include "F:\Pratice\vc98\myntlib\myntlib.h"
#pragma comment(lib,"F:/Pratice/vc98/myntlib/Release/myntlib.lib")
#pragma comment(lib,"Dnsapi.lib")

//#include "rdb.h"
#define malloc(x) HeapAlloc(GetProcessHeap(), 0, x)
#define free(x)  HeapFree(GetProcessHeap(),0,x)
ULONG Inet_Addr(LPWSTR String, LPWSTR * Terminator);
__inline unsigned long __htonl(unsigned long n)
{
    return ((n & 0xFF) << 24) | ((n & 0xFF00) << 8) | ((n & 0xFF0000) >> 8) | ((n & 0xFF000000) >> 24);
}

__inline unsigned long __ntohl(unsigned long n)
{
    return ((n & 0xFF) << 24) | ((n & 0xFF00) << 8) | ((n & 0xFF0000) >> 8) | ((n & 0xFF000000) >> 24);
}
RDB rdb;
int initrdb;
WCHAR** filters = 0;
struct tuple{
    WCHAR* q;
    WCHAR* r;
} *gsub = 0;
extern "C" BOOL __stdcall DnsFlushResolverCache();
static HKEY key;
static HANDLE hEvent;
int process();
void setup();
void cleanup();

hook_info hook;
struct DNS_Q {
    wchar_t* name;
    DWORD u1
#if NONE
        ,u2
#endif
        ,query_type;
    DWORD flags;
    DWORD us[10];
    PDNS_RECORDW records;
    DWORD us2[2];
    PIP4_ARRAY arry;    
};
typedef int (__stdcall *Query_Main_t)(void*);
typedef PDNS_RECORD (__stdcall *DnsAllocateRecord_t)(DWORD);
static DnsAllocateRecord_t DnsAllocateRecord;
Query_Main_t real_query;
//#define mdebug(X)  OutputDebugStringA(X)
//#define wdebug(X)  OutputDebugStringW(X)
#define mdebug(X)
#define wdebug(X)
bool filter(WCHAR* name) {
    if (filters) {
        WCHAR** p = filters;
        while(*p != 0) {
            if (wcsstr(name,*p))
                return true;
            p++;
        }
    }
    return false;
    if (wcsstr(name,L"youtube")  ||  wcsstr(name,L"facebook") || wcsstr(name,L"ytimg") || wcsstr(name,L"blog.") || wcsstr(name,L"blogger") || wcsstr(name,L"blogspot") ||  wcsstr(name,L"googleusercontent.com") ||  wcsstr(name,L"xmarks.com")|| wcsstr(name,L"imdb.com")
         || wcsstr(name,L"twitter.com") ) 
        return true;
    return false;
}
IP4_ARRAY freedns = {
1,0x08080808
};
int __stdcall Query_Main(DNS_Q* query){
    mdebug("Query_Main");
    wdebug(query->name);
    int isf = 0;
    /*if (wcsstr(query->name,L"encrypted.google.com") || wcsstr(query->name,L"docs.google.com") 
        || wcsstr(query->name,L"spreadsheets.google.com") || wcsstr(query->name,L"code.google.com")  || wcsstr(query->name,L"appengine.google.com") 
        || wcsstr(query->name,L"appspot.com") 
        || wcsstr(query->name,L"groups.google.com") 
        ) {
        query->name = L"www.google.com.hk";
    }

    if (wcsstr(query->name,L"www.youtube.com") ||
        wcsstr(query->name,L"ytimg.com") 
        ) {
        query->name = L"r11.nrt05s03.c.youtube.com";
    }*/
    //debug_print("Q_M: name:%S;u1:%x;u2:%x;query_type:%d;flags:%x",query->name,query->u1,query->u2,query->query_type,query->flags);
    WCHAR* orig_name = 0;
    if (gsub) {
        WCHAR* match = query->name;
        bool matched = false;
again:
        tuple* a = &gsub[0];
        while (a->q && a->r) {
            if (wcsstr(match,a->q)) {
                matched = true;
                match = a->r;
                goto again;
            }
            a++;
        }

        if (matched) {
            orig_name = query->name;
            query->name = match;
        }
    }
    if(filter(query->name)) {
        query->flags |= DNS_QUERY_USE_TCP_ONLY;
        query->arry = &freedns;
        isf = 1;
    }
    WCHAR* end = 0;
    int ret = 0;
    ULONG address = 0;
    if (orig_name && query->query_type == DNS_TYPE_A && (address =  Inet_Addr(query->name,&end) ) != INADDR_NONE) {
        PDNS_RECORDW record = (PDNS_RECORDW)DnsAllocateRecord(sizeof(DNS_A_DATA));
        record->pName = orig_name;
        record->Data.A.IpAddress = address;
        record->wType = DNS_TYPE_A;
        record->dwTtl = 86400;
        record->Flags.S.Section = DNSREC_ANSWER;
        record->Flags.S.CharSet = DnsCharSetUnicode;
        //record->Flags.S.Reserved = 0X20; //iF there is more records afterwards
        query->records = record;
    }
    else 
        ret = real_query(query);
    if (!initrdb)
        return ret;
    PDNS_RECORDW set = query->records;
    while(set){
        if (set->wType == DNS_TYPE_A) {
            WCHAR tt[56];
            //debug_print("%S %S ttl:%d;flags:(%x %x %x %x);Type:%x;wDataLength:%d",set->pName, inet_ntow(set->Data.A.IpAddress,tt),set->dwTtl,set->Flags.S.Section,set->Flags.S.Delete,set->Flags.S.CharSet,set->Flags.S.Reserved,set->wType,set->wDataLength );
            if (geoip(__ntohl(set->Data.A.IpAddress)))
                break;
            inet_ntow(set->Data.A.IpAddress,tt);
            rdb_put(&rdb,tt,query->name);
        }
        set = set->pNext;
    }
    return ret;
}
struct SOCKET_T {
    DWORD u1[2];
    SOCKET sock;
    DWORD fromlen;
    sockaddr from;
    DWORD u2[7];
// 60 byte
    DWORD seconds;
    DWORD u3[9];
    unsigned short u4;
//102 byte
    unsigned short recved;
//104 byte
    CHAR buf[1472];
};
int __stdcall Dns_RecvUdp(void*) {
    return 0;
}
WCHAR module[MAX_PATH];
HMODULE curr_module;
BOOL __stdcall DllEntry( HMODULE hinstDLL,DWORD reason,PVOID){
    if (reason == DLL_PROCESS_ATTACH) {
        DisableThreadLibraryCalls(hinstDLL);
        curr_module = hinstDLL;
        WCHAR proc[MAX_PATH];
        GetModuleFileNameW(hinstDLL,module,sizeof(module)/sizeof(module[0]));
        if (find_process_name(GetCurrentProcessId(),proc,sizeof(proc)) && wcscmp(proc,L"svchost.exe") == 0 && GetModuleHandleW(L"dnsrslvr.dll")) {
            GetModuleHandleExW(0,module,&hinstDLL);
            HMODULE dns = GetModuleHandleW(L"dnsapi.dll");
            DnsAllocateRecord = (DnsAllocateRecord_t)GetProcAddress(dns,"DnsAllocateRecord");
            if (hook_proc(GetProcAddress(dns,"Query_Main"), Query_Main,&real_query,5,&hook)) {
                initrdb = open_rdb(L"ipnames",L"d:/ipnames",&rdb);
                setup();
                OutputDebugStringA("Hook dns successfully");
            }
        }
    }
    else if (reason == DLL_PROCESS_DETACH) {
        if (unhook_proc(&hook)) {
            cleanup();
            OutputDebugStringA("Freeed");
        }
    }
    
    return TRUE;
}
DWORD __stdcall freeme(void*) {
    FreeLibraryAndExitThread(curr_module,0);
    return 1;
}

DWORD getdnscache() {
    SC_HANDLE s = OpenSCManagerW(0,0,GENERIC_READ);
    SC_HANDLE sv = OpenServiceW(s,L"dnscache",SERVICE_QUERY_STATUS);
    SERVICE_STATUS_PROCESS status = {0};
    DWORD r;
    QueryServiceStatusEx(sv,SC_STATUS_PROCESS_INFO ,(LPBYTE )&status,sizeof(status),&r);
    return status.dwProcessId;
}

#pragma comment(linker,"/export:injW=_inject@16")
extern "C" void __stdcall inject(HWND,HINSTANCE,LPWSTR, int )
{
    DWORD pid = getdnscache();
    if (pid) {
        PVOID p = GetProcAddress(GetModuleHandleW(L"kernel32.dll"),"GetModuleHandleW");
        DWORD code = 1;
        if (call_remote_par(pid,(tproc_t)p,module,wcslen(module)+sizeof(WCHAR),&code) && code == 0) {
            if (inject_me(pid,module) == 0) {
                char tmp[128];
                sprintf(tmp,"%s %d","Failed injection",GetLastError());
                OutputDebugStringA(tmp);
            }
        }
        else {
            OutputDebugStringA("Cannot call GetModuleHandleW");
        }
    }
    else {
        OutputDebugStringA("Cannot find dnscache");
    }
}
#pragma comment(linker,"/export:freeW=_free_@16")
extern "C" void __stdcall free_(HWND,HINSTANCE,LPWSTR, int )
{
    DWORD pid = getdnscache();
    if (pid) {
        call_remote(pid,freeme,0);
    }
}











DWORD WINAPI wfunc(PVOID) {
    while(1){
        WaitForSingleObject(hEvent,INFINITE);
        process();
        RegNotifyChangeKeyValue(key,0,REG_NOTIFY_CHANGE_LAST_SET,hEvent,1);
    }
}
static HANDLE bthread;
void setup() {
    RegOpenKeyExW(HKEY_LOCAL_MACHINE,L"SYSTEM\\CurrentControlSet\\Control\\ip",0,KEY_QUERY_VALUE|KEY_NOTIFY,&key);
    hEvent = CreateEvent(NULL, FALSE, FALSE, NULL);
    process();
    RegNotifyChangeKeyValue(key,0,REG_NOTIFY_CHANGE_LAST_SET,hEvent,1);
    bthread = CreateThread(0,0,wfunc,0,0,0);
}
void cleanup() {
    TerminateThread(bthread,0);
    CloseHandle(bthread);
    CloseHandle(hEvent);
    RegCloseKey(key);
    free(gsub);
    free(filters);
}

int process() {
    
    DWORD ret = 0,type =0;
    DWORD size = 0;
    DWORD count = 0, namelen = 0, valuelen = 0;
    RegQueryInfoKeyW(key,0,0,0,0,0,0,&count,&namelen,&valuelen,0,0);
    namelen += 1;
    WCHAR* name = (WCHAR*)_alloca(namelen*2);
    LPBYTE  value = (LPBYTE)_alloca(valuelen);
    tuple* subs = (tuple*)malloc(count*sizeof(tuple));
    int ii = 0;
    for (unsigned  q = 0; q < count; q++) {
        DWORD nl = namelen, vl = valuelen;
        ret = RegEnumValueW(key,q,name,&nl,0,0,value,&vl);
        if (ret == ERROR_SUCCESS && wcscmp(name,L"afilter") != 0) {
            subs[ii].q = (WCHAR*)malloc(nl*2+2);
            wcscpy(subs[ii].q,name);
            subs[ii].r =  (WCHAR*)malloc(vl+2);
            wcsncpy(subs[ii++].r,(WCHAR*)value,vl/2);
        }
    }
    subs[ii].q=0;
    free(gsub);
    gsub = subs;
    WCHAR* buf = (WCHAR*)malloc(valuelen);
    size = valuelen;
    ret = RegQueryValueExW(key,L"afilter",0,&type,(PBYTE)buf,&size);
    if (ret != ERROR_SUCCESS)
        return -1;
    WCHAR** f = (WCHAR**)malloc(size/2*sizeof(char*));
    unsigned i = 0;
    int j = 0;
    while (i < size) {
        int len = wcslen(&buf[i]);
        if (len > 0) {
            f[j++] = &buf[i];
            i += len + 1;
        }
        else {
            break;
        }
    }        
    f[j] = 0;
    if (filters && filters[0]) free(filters[0]);
    free(filters);
    filters = f;
    DnsFlushResolverCache();
    return 0;
}